import matplotlib
import numpy as np
from matplotlib import pyplot as plt

BITS = 32

def init_direction2():
    s = 1
    a = 0
    m = [0, 1]
    direction_number = [0] * BITS
    for i in range(1, s + 1):
        direction_number[i] = m[i] << BITS - i
    for i in range(s + 1, BITS):
        direction_number[i] = direction_number[i -s] ^ (direction_number[i -s] >> s)
        for k in range(1, s):
            direction_number[i - 1] = ((a >> (s - 1 - k)) & 1) * direction_number[i - k - 1]
    return direction_number

# -------------------------------------------------
# 读取方向数（这里直接写入前 3 维的示例）
direction_numbers = {
    1: [1 << (BITS - i + 1) for i in range(1, BITS)],                     # 2⁻¹, 2⁻², …
    2: init_direction2()
}
direction_numbers[1][0] = 0

def sobol_algorithm1(index: int, dim: int) -> np.ndarray:
    """返回第 index（从 0 开始）的 Sobol 点，维度 dim"""
    index <<= 1
    g = index ^ (index >> 1)
    x = np.zeros(dim)
    for d in range(dim):
        # 方向数组
        v = direction_numbers[d + 1]          # 维度从 1 开始计数
        # 逐位与运算并累加
        val = 0
        for k, vk in enumerate(v):
            if (g >> k) & 1:                 # 第 k 位为 1
                val ^= vk
        x[d] = val / 2 ** BITS                    # 归一化到 [0,1)
    return x

state = [0, 0]

def sobol_algorithm2(index: int, dim: int) -> np.ndarray:
    """返回第 index（从 0 开始）的 Sobol 点，维度 dim"""
    if index == 0:
        return np.zeros(dim)
    c = 1
    value = index - 1
    while ((value & 1) == 1) :
        value >>= 1
        c+=1

    v = np.zeros(dim)

    for i in range(dim):
        state[i] ^= direction_numbers[i+1][c]
        v[i] = state[i] /2 ** BITS

    return v


def standard_math(ax):
    # 先隐藏上、右两条边框
    ax.spines['top'].set_color('none')
    ax.spines['right'].set_color('none')

    # 将左、下两条边框移动到 (0,0) 位置
    ax.spines['left'].set_position('zero')
    ax.spines['bottom'].set_position('zero')

    # -------------------------------------------------
    # 5️⃣ 添加箭头（模拟数学教材中的“→”和“↑”）
    # -------------------------------------------------
    # 这里使用 annotate 手动在坐标轴末端绘制箭头
    arrow_style = dict(arrowstyle='->', linewidth=1.5, color='black')
    ax.annotate('', xy=(np.pi, 0), xytext=(0, 0), arrowprops=arrow_style)  # x‑轴正方向
    ax.annotate('', xy=(0, 1.5), xytext=(0, 0), arrowprops=arrow_style)  # y‑轴正方向

    # -------------------------------------------------
    # 6️⃣ 设置刻度只显示在轴的内部（数学坐标系的常规做法）
    # -------------------------------------------------
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    #ax.set_xticks(np.arange(-np.pi, np.pi + 0.1, np.pi / 2))
    #ax.set_yticks(np.arange(-1, 2, 0.5))

    # 为了美观，去掉刻度线的外延
    ax.tick_params(direction='in', length=6, width=1)

    # -------------------------------------------------
    # 7️⃣ 添加轴标签（可选）
    # -------------------------------------------------
    ax.set_xlabel(r'$x$', fontsize=14, labelpad=10)
    ax.set_ylabel(r'$y$', fontsize=14, labelpad=10, rotation=0)

    # -------------------------------------------------
    # 8️⃣ 保持等比例（确保单位长度在 x、y 方向相同）
    # -------------------------------------------------
    ax.set_aspect('equal', adjustable='box')


if __name__ == "__main__":
    # -------------------------------------------------
    # 示例：生成前 16 个 3‑维 Sobol 点
    points1 = np.vstack([sobol_algorithm1(i, 2) for i in range(16)])
    points2 = np.vstack([sobol_algorithm2(i, 2) for i in range(16)])
    print(points1 == points2)

    matplotlib.rcParams['font.family'] = 'Microsoft YaHei'  # Windows 常用
    plt.figure(figsize=(7, 7))
    plt.scatter(points1[:, 0], points1[:, 1],
                color="steelblue", s=60, edgecolor="k", alpha=0.8)

    # 为每个点标注序号（可选）
    for i, (x, y) in enumerate(points1):
        plt.text(x, y+0.01, str(i), fontsize=8,
                 ha='center', va='bottom', color='darkred')

    plt.title("Sobol 序列前 16 点（2 维）")
    plt.xlabel("Dimension 1")
    plt.ylabel("Dimension 2")
    plt.xlim(0, 1)
    plt.ylim(0, 1)
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.gca().set_aspect('equal', adjustable='box')

    plt.tight_layout()

    # standard_math(plt.gca())

    plt.show()
